# -*- coding: utf-8 -*-
"""
Created on Fri Oct  2 09:49:38 2020

@author: jyh
"""

"""
数据选择模块的代码
"""
import numpy as np

# 对于每种数据，如果长度小于 L，则进行填充
def pad_data(data, length):
    padded_data = np.zeros((data.shape[0], length), dtype=data.dtype)
    for i, row in enumerate(data):
        padded_data[i, :min(length, row.size)] = row[:min(length, row.size)]
    return padded_data

def load_data_payload(n_text_pain=0,n_text_encr=0,n_acars_pain=0,n_acars_encr=0,n_ais_pain=0,n_ais_encr=0,n_ais1_pain=0,n_ais1_encr=0,\
              n_ais4_pain=0,n_ais4_encr=0,n_http=0,n_ssh=0,n_dns=0,n_ftp=0,n_smtp=0,L=10):
    #载荷数据
    '''Text'''
    pain_text = np.load("payload\painText_dec_all.npy") #96194条，1024字节长
    Blowfishcipher_text = np.load("payload\BlowfishcipherText_dec_all.npy")#46580条，1024字节长
    AEScipher_text = np.load("payload\AEScipherText_dec_all.npy")#46580条，1024字节长
    encr_text = np.vstack((Blowfishcipher_text,AEScipher_text))#93160条，1024字节长
    
    '''ACARS'''
    pain_acars = np.load(r"payload\acars_pain_dec_9930.npy") #9930条，48字节长
    encr_acars = np.load(r"payload\acars_aes_dec_9930.npy") #9930条，48字节长
    
    '''AIS'''
    #ais各种消息混杂
    pain_ais_all = np.load(r"payload\AIS_pain_dec_51678.npy")#51678条，16字节长
    encr_ais_all = np.load(r"payload\AIS_aes_dec_51678.npy")#51678条，16字节长
    #ais消息1
    pain_ais1 = np.load(r"payload\AIS_pain_dec_xiaoxi1.npy")#150460条，16字节长
    encr_ais1 = np.load(r"payload\AIS_aes_dec_xiaoxi1.npy")#150460条，16字节长
    #ais消息4
    pain_ais4 = np.load(r"payload\AIS_pain_dec_xiaoxi4.npy")#6637条，16字节长
    encr_ais4 = np.load(r"payload\AIS_aes_dec_xiaoxi4.npy")#6637条，16字节长
    
    '''HTTP'''
    http = np.load(r"payload\http_whole_fram_dec_payload(40B).npy")#153831条，40字节长
    
    '''SSH'''
    ssh = np.load(r"payload\ssh_dec_payload(40B).npy")#11608条，40字节长
    
    '''DNS'''
    dns = np.load(r"payload\dns_maccdc_dec.npy")#16137条，16字节长
    '''FTP'''
    ftp = np.load(r"payload\ftp_maccdc_dec.npy")#1569条，16字节长
    '''SMTP'''
    smtp = np.load(r"payload\smtp_maccdc_dec.npy")#181条，16字节长

    # 对每个数据集进行填充
    pain_text = pad_data(pain_text, L)
    encr_text = pad_data(encr_text, L)
    pain_acars = pad_data(pain_acars, L)
    encr_acars = pad_data(encr_acars, L)
    pain_ais_all = pad_data(pain_ais_all, L)
    encr_ais_all = pad_data(encr_ais_all, L)
    pain_ais1 = pad_data(pain_ais1, L)
    encr_ais1 = pad_data(encr_ais1, L)
    pain_ais4 = pad_data(pain_ais4, L)
    encr_ais4 = pad_data(encr_ais4, L)
    http = pad_data(http, L)
    ssh = pad_data(ssh, L)
    dns = pad_data(dns, L)
    ftp = pad_data(ftp, L)
    smtp = pad_data(smtp, L)
    
    pain_data = np.vstack((pain_text[0:n_text_pain,0:L],pain_acars[0:n_acars_pain,0:L],pain_ais_all[0:n_ais_pain,0:L],pain_ais1[0:n_ais1_pain,0:L],\
                      pain_ais4[0:n_ais4_pain,0:L],http[0:n_http,0:L],dns[0:n_dns,0:L],ftp[0:n_ftp,0:L],smtp[0:n_smtp,0:L]))
    encr_data = np.vstack((encr_text[0:n_text_encr,0:L],encr_acars[0:n_acars_encr,0:L],encr_ais_all[0:n_ais_encr,0:L],encr_ais1[0:n_ais1_encr,0:L],\
                      encr_ais4[0:n_ais4_encr,0:L],ssh[0:n_ssh,0:L]))
    return pain_data,encr_data

if __name__ == "__main__":
    pain,encr = load_data_payload(0,0,10,0,L=40)